#!/usr/bin/env python
"""
初始化百家乐系统PostgreSQL数据库表
"""
import os
import sys
import psycopg2
from dotenv import load_dotenv
import time
import argparse

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def drop_tables(db_config, confirm=False):
    """删除所有现有表"""
    if not confirm:
        confirmation = input("确定要删除所有现有表吗？这将清除所有数据！(y/n): ")
        if confirmation.lower() != 'y':
            print("操作已取消")
            return False

    conn = None
    cursor = None
    
    try:
        # 连接到PostgreSQL服务器
        print(f"正在连接到PostgreSQL数据库: {db_config['host']}:{db_config['port']}/{db_config['database']}")
        conn = psycopg2.connect(**db_config)
        cursor = conn.cursor()
        
        # 开启事务
        conn.autocommit = False
        
        # 按照依赖关系顺序删除表（先删除引用表，再删除被引用表）
        print("正在删除现有表...")
        
        # 1. 删除 simulation_rounds 表 (依赖于 simulations)
        cursor.execute("DROP TABLE IF EXISTS simulation_rounds CASCADE")
        print("- simulation_rounds 表已删除")
        
        # 2. 删除 decks 表 (依赖于 simulations)
        cursor.execute("DROP TABLE IF EXISTS decks CASCADE")
        print("- decks 表已删除")
        
        # 3. 删除 simulations 表
        cursor.execute("DROP TABLE IF EXISTS simulations CASCADE")
        print("- simulations 表已删除")
        
        # 4. 删除 rounds 表
        cursor.execute("DROP TABLE IF EXISTS rounds CASCADE")
        print("- rounds 表已删除")
        
        # 5. 删除 bet_strategy_tables 表
        cursor.execute("DROP TABLE IF EXISTS bet_strategy_tables CASCADE")
        print("- bet_strategy_tables 表已删除")
        
        # 提交事务
        conn.commit()
        print("所有表删除成功！")
        return True
    
    except Exception as e:
        print(f"删除表错误: {str(e)}")
        if conn:
            conn.rollback()
        return False
    
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()

def init_database(db_config):
    """初始化数据库表结构"""
    conn = None
    cursor = None
    
    try:
        # 连接到PostgreSQL服务器
        print(f"正在连接到PostgreSQL数据库: {db_config['host']}:{db_config['port']}/{db_config['database']}")
        conn = psycopg2.connect(**db_config)
        cursor = conn.cursor()
        
        # 开启事务
        conn.autocommit = False
        
        # 1. 创建普通游戏记录表
        print("创建游戏记录表 (rounds)...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS rounds (
                id SERIAL PRIMARY KEY,
                timestamp TEXT,
                bet_choice TEXT,
                bet_amount REAL,
                result TEXT,
                win_or_lose INTEGER,
                payout REAL,
                player_hand TEXT,
                banker_hand TEXT,
                player_points INTEGER,
                banker_points INTEGER,
                main_choice TEXT
            )
        ''')
        
        # 2. 创建模拟记录表
        print("创建模拟记录表 (simulations)...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS simulations (
                id SERIAL PRIMARY KEY,
                timestamp TEXT,
                strategy_name TEXT,
                total_rounds INTEGER,
                win_rounds INTEGER,
                lose_rounds INTEGER,
                net_profit REAL,
                deck_seed TEXT,
                parameters TEXT,
                main_choice TEXT
            )
        ''')
        
        # 3. 创建模拟详细记录表
        print("创建模拟详细记录表 (simulation_rounds)...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS simulation_rounds (
                id SERIAL PRIMARY KEY,
                timestamp TEXT,
                simulation_id INTEGER REFERENCES simulations(id),
                round_id INTEGER,
                bet_choice TEXT,
                bet_amount REAL,
                result TEXT,
                win_or_lose INTEGER,
                payout REAL,
                player_hand TEXT,
                banker_hand TEXT,
                player_points INTEGER,
                banker_points INTEGER,
                strategy_note TEXT,
                main_choice TEXT
            )
        ''')
        
        # 4. 创建牌组记录表
        print("创建牌组记录表 (decks)...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS decks (
                id SERIAL PRIMARY KEY,
                timestamp TEXT,
                simulation_id INTEGER REFERENCES simulations(id),
                deck_data TEXT,
                original_size INTEGER,
                deck_seed INTEGER
            )
        ''')
        
        # 5. 创建下注策略表
        print("创建下注策略表 (bet_strategy_tables)...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS bet_strategy_tables (
                id SERIAL PRIMARY KEY,
                name TEXT UNIQUE,
                description TEXT,
                strategy_content TEXT DEFAULT '',
                created_at TEXT,
                updated_at TEXT,
                is_default BOOLEAN DEFAULT FALSE
            )
        ''')
        
        # 提交事务
        conn.commit()
        print("数据库表创建成功！")
        
        # 导入默认下注策略表
        import_default_strategy_table(db_config)
        
        return True
    
    except Exception as e:
        print(f"数据库初始化错误: {str(e)}")
        if conn:
            conn.rollback()
        return False
    
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()

def import_default_strategy_table(db_config):
    """导入默认下注策略表"""
    try:
        # 从config导入默认策略表
        conn = psycopg2.connect(**db_config)
        cursor = conn.cursor()
        
        # 先检查是否已有默认策略表
        cursor.execute("SELECT COUNT(*) FROM bet_strategy_tables WHERE name = '默认策略'")
        count = cursor.fetchone()[0]
        
        if count == 0:
            print("导入默认下注策略表...")
            # 准备默认策略内容
            strategy_content = """# 默认策略规则
# 格式说明：
# +1000:1000 表示上局赢了1000，本局下注1000
# -1000:1100 表示上局输了1000，本局下注1100
+1000:1000
-1000:1100
-1100:1300
-1300:1700
-1700:1000
"""
            created_at = updated_at = time.strftime('%Y-%m-%d %H:%M:%S')
            
            # 插入默认策略表
            cursor.execute('''
                INSERT INTO bet_strategy_tables 
                (name, description, strategy_content, created_at, updated_at, is_default)
                VALUES (%s, %s, %s, %s, %s, %s)
            ''', (
                '默认策略', 
                '系统默认的下注策略表，根据历史输赢记录调整下注金额', 
                strategy_content,
                created_at,
                updated_at,
                True
            ))
            conn.commit()
            print("默认下注策略表导入成功！")
        else:
            print("默认下注策略表已存在，跳过导入")
            
    except Exception as e:
        print(f"导入默认下注策略表错误: {str(e)}")
        if conn:
            conn.rollback()
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()

def check_database_exists(db_config, db_name):
    """检查数据库是否存在，不存在则创建"""
    # 连接到默认的postgres数据库
    conn_params = {
        'host': db_config['host'],
        'port': db_config['port'],
        'user': db_config['user'],
        'password': db_config['password'],
        'database': 'postgres'  # 连接到默认postgres数据库
    }
    
    conn = None
    cursor = None
    
    try:
        # 连接到postgres数据库
        conn = psycopg2.connect(**conn_params)
        conn.autocommit = True  # 自动提交
        cursor = conn.cursor()
        
        # 检查数据库是否存在
        cursor.execute("SELECT 1 FROM pg_database WHERE datname = %s", (db_name,))
        exists = cursor.fetchone()
        
        if not exists:
            print(f"数据库 '{db_name}' 不存在，正在创建...")
            # 创建数据库
            cursor.execute(f"CREATE DATABASE {db_name}")
            print(f"数据库 '{db_name}' 创建成功！")
            return True
        else:
            print(f"数据库 '{db_name}' 已存在")
            return True
    
    except Exception as e:
        print(f"检查/创建数据库错误: {str(e)}")
        return False
    
    finally:
        if cursor:
            cursor.close()
        if conn:
            conn.close()

def main():
    """主函数"""
    print("="*50)
    print("百家乐系统数据库初始化工具")
    print("="*50)
    
    # 命令行参数解析
    parser = argparse.ArgumentParser(description='初始化百家乐系统数据库')
    parser.add_argument('-d', '--drop', action='store_true', help='删除所有现有表并重新创建')
    parser.add_argument('-f', '--force', action='store_true', help='强制执行，不询问确认')
    args = parser.parse_args()
    
    # 加载环境变量
    load_dotenv()
    
    # 数据库连接参数
    db_config = {
        'host': os.getenv('DB_HOST', 'localhost'),
        'port': int(os.getenv('DB_PORT', 5432)),
        'user': os.getenv('DB_USER', 'postgres'),
        'password': os.getenv('DB_PASSWORD', 'Postgre,.1'),
        'database': os.getenv('DB_NAME', 'baijiale')
    }
    
    db_name = db_config['database']
    
    # 检查数据库是否存在
    if check_database_exists(db_config, db_name):
        # 如果指定删除表选项
        if args.drop:
            if drop_tables(db_config, args.force):
                # 重新初始化表结构
                if init_database(db_config):
                    print("\n数据库重新初始化完成！")
                else:
                    print("\n数据库表初始化失败！")
            else:
                print("\n删除表操作未完成！")
        else:
            # 正常初始化表结构
            if init_database(db_config):
                print("\n数据库初始化完成！")
            else:
                print("\n数据库表初始化失败！")
    else:
        print("\n数据库检查/创建失败！")

if __name__ == "__main__":
    main() 